from torch.utils.data import Dataset
from pathlib import Path
import pandas as pd
import cv2
import numpy as np

class MNISTDataset(Dataset):
    def __init__(self, set_path, index_path, use_cache=False):
        super().__init__()

        base_path = Path(set_path)
        self.used_cache = use_cache

        index_data = pd.read_csv(index_path, header=None, sep="   ", engine="python")
        self.index_map = dict()
        for i in range(index_data.shape[0]):
            self.index_map[str(index_data.iloc[i, 0])] = int(index_data.iloc[i, 1])
       
        if use_cache:
            self.files = [self._deal_one(it) for it in base_path.iterdir()]
        else:
            self.files = [it for it in base_path.iterdir()]

    def _deal_one(self, it):
        img = cv2.imread(str(it), cv2.IMREAD_GRAYSCALE) / 255.0
        return img.reshape(-1, 28, 28).astype(np.float32), img.flatten().astype(np.float32)
    

    def __len__(self):
        return len(self.files)
    

    def __getitem__(self, idx):
        if not self.used_cache:
            return self._deal_one(self.files[idx])
        
        return self.files[idx]